package com.jy.print.javaprint.utils;

import org.eclipse.swt.graphics.ImageData;
import org.eclipse.swt.graphics.PaletteData;
import org.eclipse.swt.graphics.RGB;
import org.opencv.core.CvType;
import org.opencv.core.Mat;

import java.awt.image.BufferedImage;
import java.awt.image.ColorModel;
import java.awt.image.DirectColorModel;
import java.awt.image.WritableRaster;

public class BufImgToMat {

    public static ImageData getImageData(BufferedImage bufferedImage) {
        ColorModel colorModel = bufferedImage.getColorModel();
        PaletteData palette = null;
        if (colorModel instanceof DirectColorModel) {
            palette = new PaletteData(((DirectColorModel) colorModel).getRedMask(), ((DirectColorModel) colorModel).getGreenMask(), ((DirectColorModel) colorModel).getBlueMask());
        } else {
            palette = new PaletteData(0x0000FF, 0x00FF00, 0xFF0000);
        }
        ImageData data = new ImageData(bufferedImage.getWidth(), bufferedImage.getHeight(), colorModel.getPixelSize(), palette);
        WritableRaster raster = bufferedImage.getRaster();
        int[] pixelArray = new int[3];
        for (int y = 0; y < data.height; y++) {
            for (int x = 0; x < data.width; x++) {
                raster.getPixel(x, y, pixelArray);
                int pixel = palette.getPixel(new RGB(pixelArray[0], pixelArray[1], pixelArray[2]));
                data.setPixel(x, y, pixel);
            }
        }
        return data;
    }

    BufferedImage	original;

    public BufImgToMat(BufferedImage image) {
        original = image;
    }

    public Mat getMat() {
        Mat out;
        byte[] data;
        int r, g, b;
        int w = original.getWidth();
        int h = original.getHeight();
        if (original.getType() == BufferedImage.TYPE_INT_RGB) {
            out = new Mat(h, w, CvType.CV_8UC3);
            data = new byte[w * h * (int) out.elemSize()];
            int[] dataBuff = original.getRGB(0, 0, w, h, null, 0, w);
            for (int i = 0; i < dataBuff.length; i++) {
                data[i * 3] = (byte) ((dataBuff[i] >> 16) & 0xFF);
                data[i * 3 + 1] = (byte) ((dataBuff[i] >> 8) & 0xFF);
                data[i * 3 + 2] = (byte) ((dataBuff[i] >> 0) & 0xFF);
            }
        } else {
            out = new Mat(h, w, CvType.CV_8UC1);
            data = new byte[w * h * (int) out.elemSize()];
            int[] dataBuff = original.getRGB(0, 0, w, h, null, 0, w);
            for (int i = 0; i < dataBuff.length; i++) {
                r = (byte) ((dataBuff[i] >> 16) & 0xFF);
                g = (byte) ((dataBuff[i] >> 8) & 0xFF);
                b = (byte) ((dataBuff[i] >> 0) & 0xFF);
                data[i] = (byte) ((0.21 * r) + (0.71 * g) + (0.07 * b));
            }

        }
        out.put(0, 0, data);
        return out;
    }

}